%{
#include <kuai/sim/simjob.hpp>
#include <kuai/sim/md.hpp>

%}

namespace kuai { 

	class SimulationJob
	{
	public:
		SimulationJob();
		virtual ~SimulationJob();

	public:
		int random_seed;
		int getCPU();
		void setCPU(int i);
		
		%pythoncode %{
			cpu = property(getCPU, setCPU)
		%}
	};
	
	%extend SimulationJob {
		PyObject* getE(PyObject* model, PyObject* efunc) {
			EnergySet eset;
			if (!convertToEnergySet(model, efunc, eset)) {
				return NULL;
			}
			std::pair<PyObject*, RealNumber*> e = newRealArray(eset.getRank(), 0.0);
			if (e.first != NULL) {
				$self->addE(eset, e.second);
			}
			return e.first;
		}
		PyObject* getEG(PyObject* model, PyObject* efunc) {
			EnergySet eset;
			if (convertToEnergySet(model, efunc, eset)) {
				std::pair<PyObject*, RealNumber*> e = newRealArray(eset.getRank(), 0.0);
				if (e.first != NULL) {
					std::pair<PyObject*, XYZ*> g = newXYZArray(eset.nAtoms, true);
					if (g.first != NULL) {
						$self->addEG(eset, e.second, g.second);
						PyObject* result = PyTuple_Pack(2, e.first, g.first);
						Py_DECREF(e.first);
						Py_DECREF(e.second);
						return result;
					}
				}
			}
			return NULL;
		
		}
		PyObject* getEGH(PyObject* model, PyObject* efunc) {
			EnergySet eset;
			if (convertToEnergySet(model, efunc, eset)) {
				std::pair<PyObject*, RealNumber*> e = newRealArray(eset.getRank(), 0.0);
				if (e.first != NULL) {
					std::pair<PyObject*, XYZ*> g = newXYZArray(eset.nAtoms, true);
					if (g.first != NULL) {
						std::pair<PyObject*, RealNumber*> h = newRealArray(9 * eset.nAtoms * eset.nAtoms, 0.0);
						if (h.first != NULL) {
							std::vector<RealNumber*> hess(3 * eset.nAtoms);
							for (Index i = 0; i < hess.size(); ++i) {
								hess[i] = h.second + i * 3 * eset.nAtoms;
							}
							$self->addEGH(eset, e.second, g.second, &hess[0]);
							PyObject* result = PyTuple_Pack(3, e.first, g.first, h.first);
							Py_DECREF(e.first);
							Py_DECREF(e.second);
							return result;
						}
					}
				}
			}
			return NULL;
		
		}

		PyObject* getHX(PyObject* model, PyObject* efunc, PyObject* x) {
			EnergySet eset;
			if (convertToEnergySet(model, efunc, eset)) {
				
				std::pair<PyObject*, RealNumber*> result = newRealArray(eset.nAtoms*3, 0.0);
				if (result.first != NULL) {
					std::pair<RealNumber*, size_t>  xx = getRealArray(x);
					if (xx.second == eset.nAtoms*3) {
						$self->addHX(eset,  xx.first, result.second);
						return result.first;
					}
				}
			}
			return NULL;
		
		}
	};
	
	class MolecularDynamicJob
		: public SimulationJob
	{ };
	
	%extend MolecularDynamicJob {
	
		void run_md(PyObject* model, PyObject* efunc, PyObject* controls, size_t nstep, RealNumber timestep=1e-15) {
			EnergySet eset;
			if (!convertToEnergySet(model, efunc, eset)) {
				return;
			}

			if (!MolecularDynamicControlProxy::prepare(NULL, model, eset, *($self->md_control))) {
				return;
			}
			ControlSet cset;
			std::vector<PythonControl*> pcontrols;
			if (!convertToControlSet(model, eset, controls, cset, pcontrols)) {
				for (size_t i = 0; i < pcontrols.size(); ++i) {
					delete pcontrols[i];
				}
				return;
			}
			$self->run_md(eset, cset, nstep, timestep);
			for (size_t i = 0; i < pcontrols.size(); ++i) {
				delete pcontrols[i];
			}
		}
	};

}

